use std::collections::HashSet;
use std::collections::VecDeque;

fn parse_input() -> Vec<Vec<u32>> {
    let input = std::fs::read_to_string("./day11/input.txt").unwrap();
    input.lines().map(|line| -> Vec<u32> {
        line.chars().map(|ch| ch.to_digit(10).unwrap() ).collect()
    }).collect()
}

#[allow(unused)]
fn part1(mut energy_array: Vec<Vec<u32>>) {
    type Coord = (i32, i32);

    fn update(energy_array: &mut Vec<Vec<u32>>) -> usize {
        let row = energy_array.len() as i32;
        let col = energy_array[0].len() as i32;

        let in_range = |&(x, y): &Coord| -> bool {
            x >= 0 && x < col && y >= 0 && y < row
        };

        let adjacents = |&(x, y): &Coord| -> Vec<Coord> {
            [
                (x - 1, y - 1), (x    , y - 1), (x + 1, y - 1),
                (x - 1, y    ),                 (x + 1, y    ),
                (x - 1, y + 1), (x    , y + 1), (x + 1, y + 1),
            ].into_iter().filter(|coord| in_range(coord)).collect()
        };

        // Increase all energy
        let mut queue = VecDeque::new();
        let mut flashed_pos = HashSet::new();
        
        for y in 0..row {
            for x in 0..col {
                if energy_array[y as usize][x as usize] == 9 {
                    queue.push_back((x, y));
                    flashed_pos.insert((x, y));
                } else {
                    energy_array[y as usize][x as usize] += 1;
                }
            }
        }
                
        loop {
            match queue.pop_front() {
                None => break,
                Some(pos)  => {
                    // increase adj                    
                    for adj_pos in adjacents(&pos) {
                        if !flashed_pos.contains(&adj_pos) {
                            let (x, y) = adj_pos;
                            if energy_array[y as usize][x as usize] == 9 {
                                queue.push_back((x, y));
                                flashed_pos.insert((x, y));
                            } else {
                                energy_array[y as usize][x as usize] += 1;
                            }
                        }
                    }
                }
            }
        }

        let count = flashed_pos.len();

        // flashed to 0
        for (x, y) in flashed_pos {
            energy_array[y as usize][x as usize] = 0;
        }

        count
    }

    let mut result = 0;
    for _ in 0..100 {
        result += update(&mut energy_array);
    }

    println!("Part1: {}", result);
}

#[allow(unused)]
fn part2(mut energy_array: Vec<Vec<u32>>) {
    type Coord = (i32, i32);

    fn update(energy_array: &mut Vec<Vec<u32>>) -> usize {
        let row = energy_array.len() as i32;
        let col = energy_array[0].len() as i32;

        let in_range = |&(x, y): &Coord| -> bool {
            x >= 0 && x < col && y >= 0 && y < row
        };

        let adjacents = |&(x, y): &Coord| -> Vec<Coord> {
            [
                (x - 1, y - 1), (x    , y - 1), (x + 1, y - 1),
                (x - 1, y    ),                 (x + 1, y    ),
                (x - 1, y + 1), (x    , y + 1), (x + 1, y + 1),
            ].into_iter().filter(|coord| in_range(coord)).collect()
        };

        // Increase all energy
        let mut queue = VecDeque::new();
        let mut flashed_pos = HashSet::new();
        
        for y in 0..row {
            for x in 0..col {
                if energy_array[y as usize][x as usize] == 9 {
                    queue.push_back((x, y));
                    flashed_pos.insert((x, y));
                } else {
                    energy_array[y as usize][x as usize] += 1;
                }
            }
        }
                
        loop {
            match queue.pop_front() {
                None => break,
                Some(pos)  => {
                    // increase adj                    
                    for adj_pos in adjacents(&pos) {
                        if !flashed_pos.contains(&adj_pos) {
                            let (x, y) = adj_pos;
                            if energy_array[y as usize][x as usize] == 9 {
                                queue.push_back((x, y));
                                flashed_pos.insert((x, y));
                            } else {
                                energy_array[y as usize][x as usize] += 1;
                            }
                        }
                    }
                }
            }
        }

        let count = flashed_pos.len();

        // flashed to 0
        for (x, y) in flashed_pos {
            energy_array[y as usize][x as usize] = 0;
        }

        count
    }

    let octopuse_count = energy_array.len() * energy_array[0].len();
    for time in 1.. {
        let flash_count = update(&mut energy_array);
        if flash_count == octopuse_count {
            println!("Part2: {}", time);
            break;
        }
    }
}


fn main() {
    let energy_array = parse_input();
    part1(energy_array.clone());
    part2(energy_array);
}
